"""
震动衰减模型

实现了不同类型的震动衰减模型，包括：
- 指数衰减模型
- 线性衰减模型
- 幂律衰减模型
"""

import numpy as np
from typing import Union, Literal
from loguru import logger


class DecayModel:
    """震动衰减模型基类"""
    
    def __init__(self, decay_type: Literal["exponential", "linear", "power_law"] = "exponential"):
        """
        初始化衰减模型
        
        Args:
            decay_type: 衰减类型，支持指数、线性、幂律衰减
        """
        self.decay_type = decay_type
        logger.info(f"初始化{decay_type}衰减模型")
    
    def calculate_amplitude(self, 
                          initial_amplitude: float, 
                          distance: Union[float, np.ndarray], 
                          attenuation_coefficient: float) -> Union[float, np.ndarray]:
        """
        计算衰减后的振幅
        
        Args:
            initial_amplitude: 初始振幅
            distance: 传播距离
            attenuation_coefficient: 衰减系数
            
        Returns:
            衰减后的振幅
        """
        if self.decay_type == "exponential":
            return self._exponential_decay(initial_amplitude, distance, attenuation_coefficient)
        elif self.decay_type == "linear":
            return self._linear_decay(initial_amplitude, distance, attenuation_coefficient)
        elif self.decay_type == "power_law":
            return self._power_law_decay(initial_amplitude, distance, attenuation_coefficient)
        else:
            raise ValueError(f"不支持的衰减类型: {self.decay_type}")
    
    def _exponential_decay(self, 
                          initial_amplitude: float, 
                          distance: Union[float, np.ndarray], 
                          attenuation_coefficient: float) -> Union[float, np.ndarray]:
        """
        指数衰减模型: A = A0 * e^(-αd)
        
        Args:
            initial_amplitude: 初始振幅 A0
            distance: 传播距离 d
            attenuation_coefficient: 衰减系数 α
            
        Returns:
            衰减后的振幅
        """
        return initial_amplitude * np.exp(-attenuation_coefficient * distance)
    
    def _linear_decay(self, 
                     initial_amplitude: float, 
                     distance: Union[float, np.ndarray], 
                     attenuation_coefficient: float) -> Union[float, np.ndarray]:
        """
        线性衰减模型: A = A0 * (1 - αd)
        
        Args:
            initial_amplitude: 初始振幅 A0
            distance: 传播距离 d
            attenuation_coefficient: 衰减系数 α
            
        Returns:
            衰减后的振幅
        """
        decayed = initial_amplitude * (1 - attenuation_coefficient * distance)
        # 确保振幅不为负
        return np.maximum(decayed, 0)
    
    def _power_law_decay(self, 
                        initial_amplitude: float, 
                        distance: Union[float, np.ndarray], 
                        attenuation_coefficient: float) -> Union[float, np.ndarray]:
        """
        幂律衰减模型: A = A0 / (1 + αd)^n
        
        Args:
            initial_amplitude: 初始振幅 A0
            distance: 传播距离 d
            attenuation_coefficient: 衰减系数 α
            
        Returns:
            衰减后的振幅
        """
        power = 2  # 幂律指数，可根据需要调整
        return initial_amplitude / (1 + attenuation_coefficient * distance) ** power
    
    def get_decay_curve(self, 
                       initial_amplitude: float, 
                       max_distance: float, 
                       attenuation_coefficient: float, 
                       num_points: int = 100) -> tuple[np.ndarray, np.ndarray]:
        """
        获取衰减曲线数据
        
        Args:
            initial_amplitude: 初始振幅
            max_distance: 最大距离
            attenuation_coefficient: 衰减系数
            num_points: 采样点数
            
        Returns:
            (距离数组, 振幅数组)
        """
        distances = np.linspace(0, max_distance, num_points)
        amplitudes = self.calculate_amplitude(initial_amplitude, distances, attenuation_coefficient)
        return distances, amplitudes
    
    def plot_decay_curve(self, 
                        initial_amplitude: float, 
                        max_distance: float, 
                        attenuation_coefficient: float,
                        title: str = None):
        """
        绘制衰减曲线
        
        Args:
            initial_amplitude: 初始振幅
            max_distance: 最大距离
            attenuation_coefficient: 衰减系数
            title: 图表标题
        """
        try:
            import matplotlib.pyplot as plt
            
            distances, amplitudes = self.get_decay_curve(
                initial_amplitude, max_distance, attenuation_coefficient
            )
            
            plt.figure(figsize=(10, 6))
            plt.plot(distances, amplitudes, 'b-', linewidth=2, label=f'{self.decay_type}衰减')
            plt.xlabel('传播距离 (m)')
            plt.ylabel('振幅')
            plt.title(title or f'{self.decay_type}衰减模型')
            plt.grid(True, alpha=0.3)
            plt.legend()
            plt.show()
            
        except ImportError:
            logger.warning("matplotlib未安装，无法绘制衰减曲线")
    
    def __str__(self) -> str:
        return f"DecayModel(type={self.decay_type})"
    
    def __repr__(self) -> str:
        return self.__str__() 